package me.piebridge.douban;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.KeyFactory;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.zip.ZipException;
import java.util.zip.ZipFile;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.json.JSONObject;
import org.json.JSONTokener;

import sun.misc.BASE64Decoder;

@SuppressWarnings("restriction")
public class Crack {

	protected static final String PREFXML = "com.douban.book.reader_preferences.xml";

	protected static Cipher getRSACipher(String privateKey) throws IOException, GeneralSecurityException {
		if (privateKey == null) {
			return null;
		}
		byte[] bytes = decodeBase64(privateKey);
		Cipher rsa = Cipher.getInstance("RSA");
		rsa.init(Cipher.DECRYPT_MODE, KeyFactory.getInstance("RSA").generatePrivate(new PKCS8EncodedKeySpec(bytes)));
		return rsa;
	}

	protected static Cipher getAESCipher(Cipher rsa, String key, String iv) throws IOException, GeneralSecurityException {
		SecretKeySpec mKey = new SecretKeySpec(rsa.doFinal(decodeHex(key)), "AES");
		IvParameterSpec mIv = new IvParameterSpec(rsa.doFinal(decodeHex(iv)), 0, 16);
		Cipher aes = Cipher.getInstance("AES/CBC/PKCS5Padding");
		aes.init(Cipher.DECRYPT_MODE, mKey, mIv);
		return aes;
	}

	protected static byte[] decodeBase64(String encoded) throws IOException {
		if (encoded == null) {
			return null;
		}
		return new BASE64Decoder().decodeBuffer(encoded);
	}

	protected static byte[] decodeHex(String hex) throws IOException {
		if (hex == null) {
			return null;
		}
		int length = hex.length();
		if (length == 0 || length % 2 != 0) {
			throw new IOException();
		}
		byte[] bytes = new byte[length / 2];
		for (int i = 0; i < length; i += 2) {
			bytes[i / 2] = (byte) ((Character.digit(hex.charAt(i), 16) << 4) + Character.digit(hex.charAt(i + 1), 16));
		}
		return bytes;
	}

	protected static String getPrivateKey(File path) throws IOException, FileNotFoundException {
		String magic = "<string name=\"private_key\">";
		BufferedReader bis = new BufferedReader(new FileReader(path));
		try {
			String line;
			while ((line = bis.readLine()) != null) {
				int begin = line.indexOf(magic);
				if (begin != -1) {
					int end = line.indexOf("</string>", begin + 1);
					if (end != -1) {
						return line.substring(begin + magic.length(), end);
					}
				}
			}
		} catch (IOException e) {
		} finally {
			bis.close();
		}
		return null;
	}

	private static Cipher setDefaultRSA(File directory) {
		try {
			File file;
			if (directory == null) {
				file = new File(PREFXML);
			} else {
				file = new File(directory, PREFXML);
			}
			return getRSACipher(getPrivateKey(file));
		} catch (IOException e) {
			return null;
		} catch (GeneralSecurityException e) {
			throw new Error(e);
		}
	}

	public static void decodeDirectory(Cipher rsa, File directory) throws IOException, GeneralSecurityException {
		if (rsa == null) {
			rsa = setDefaultRSA(directory);
		}
		if (rsa == null) {
			rsa = setDefaultRSA(null);
		}
		if (rsa == null) {
			System.err.println("[error] cannot find " + PREFXML);
			return;
		}
		File d;
		File m;
		if (directory != null) {
			d = new File(directory, "d");
			m = new File(directory, "m");
		} else {
			d = new File("d");
			m = new File("m");
		}
		if (m.isFile() && d.isFile()) {
			JSONObject manifest = new JSONObject(new JSONTokener(Files.readAsString(m)));
			String key = manifest.optString("key");
			String iv = manifest.optString("iv");
			Cipher aes = getAESCipher(rsa, key, iv);
			try {
				new DoubanEpub(aes, new ZipFile(d), manifest).convert();
			} catch (ZipException e) {
				System.err.println("[error] cannot decrypt " + d.getAbsolutePath());
			}
		}
	}

	public static void main(String... args) throws GeneralSecurityException, IOException {
		Cipher rsa = null;
		for (String arg : args) {
			File file = new File(arg);
			if (file.isFile()) {
				Cipher cipher = getRSACipher(getPrivateKey(file));
				if (cipher != null) {
					rsa = cipher;
				}
			} else if (file.isDirectory()) {
				decodeDirectory(rsa, file);
				for (File d : file.listFiles()) {
					if (d.isDirectory()) {
						decodeDirectory(rsa, d);
					}
				}
			}
		}
	}

}
